import socket
import base64
import hashlib
import struct

# WebSocket 握手常量
GUID = '258EAFA5-E914-47DA-95CA-C5AB0DC85B11'

# WebSocket 服务器详情
HOST = '127.0.0.1'
PORT = 8765

# 创建握手请求
def create_handshake_request(host, port):
    key = base64.b64encode(hashlib.sha1((str(socket.gethostname()) + GUID).encode()).digest()).decode()
    return (f"GET / HTTP/1.1\r\n"
            f"Host: {host}:{port}\r\n"
            f"Upgrade: websocket\r\n"
            f"Connection: Upgrade\r\n"
            f"Sec-WebSocket-Key: {key}\r\n"
            f"Sec-WebSocket-Version: 13\r\n"
            f"Origin: http://{host}\r\n\r\n"), key

# 验证握手响应
def validate_handshake_response(response, key):
    headers = response.split('\r\n')
    for header in headers:
        if 'Sec-WebSocket-Accept' in header:
            accept_value = header.split(':')[1].strip()
            expected_accept_value = base64.b64encode(hashlib.sha1((key + GUID).encode()).digest()).decode()
            return accept_value == expected_accept_value
    return False

# 创建 WebSocket 帧
def create_websocket_frame(data):
    data = data.encode('utf-8')
    length = len(data)
    if length <= 125:
        frame = struct.pack('B', 0x81) + struct.pack('B', length) + data
    elif length >= 126 and length <= 65535:
        frame = struct.pack('B', 0x81) + struct.pack('B', 126) + struct.pack('>H', length) + data
    else:
        frame = struct.pack('B', 0x81) + struct.pack('B', 127) + struct.pack('>Q', length) + data
    return frame

# 解码 WebSocket 帧
def decode_websocket_frame(frame):
    length = frame[1] & 127
    if length == 126:
        length = struct.unpack('>H', frame[2:4])[0]
        masks = frame[4:8]
        data = frame[8:]
    elif length == 127:
        length = struct.unpack('>Q', frame[2:10])[0]
        masks = frame[10:14]
        data = frame[14:]
    else:
        masks = frame[2:6]
        data = frame[6:]

    decoded = bytearray()
    for i in range(len(data)):
        decoded.append(data[i] ^ masks[i % 4])
    return decoded.decode('utf-8')

# 主函数
def main():
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    try:
        print(f"连接到 {HOST}:{PORT}...")
        sock.connect((HOST, PORT))
    except Exception as e:
        print(f"无法连接到 {HOST}:{PORT}: {e}")
        return

    handshake_request, key = create_handshake_request(HOST, PORT)
    print("握手请求:\n", handshake_request)
    sock.send(handshake_request.encode())

    response = sock.recv(1024).decode()
    print("握手响应:", response)
    if not validate_handshake_response(response, key):
        print("连接到 WebSocket 服务器失败")
        sock.close()
        return

    print("连接到 WebSocket 服务器成功")

    send_data = input("你的名字是什么: ")
    frame = create_websocket_frame(send_data)
    sock.send(frame)
    print('>>> %s' % send_data)

    frame = sock.recv(1024)
    recv_data = decode_websocket_frame(frame)
    print('<<< %s' % recv_data)

    sock.close()
    print("连接关闭")

if __name__ == "__main__":
    main()
